{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Chat-Haruhi-Suzumiya/blob/main/notebook/%E6%96%B0%E5%B0%8F%E8%AF%B4%E6%8A%BD%E5%8F%96_release.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "afxphm5ZXUdY"
      },
      "source": [
        "# 环境"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kxt1oviOXWcf"
      },
      "source": [
        "本脚本只需要在CPU环境运行\n",
        "\n",
        "如果你不在colab运行，你需要一个openai不封禁的ip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Cc9nVmSHXWHq",
        "outputId": "db01816e-447d-4e13-ad02-cb95088638f7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m10.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.5/76.5 kB\u001b[0m \u001b[31m6.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.7/1.7 MB\u001b[0m \u001b[31m18.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m6.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for google-search-results (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.5/76.5 kB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "!pip -q install langchain openai google-search-results tiktoken\n",
        "!pip -q install kor markdownify\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ysEvTErI8_xb"
      },
      "source": [
        "# 新小说提取"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O9DMlJxQwWCA"
      },
      "source": [
        "Implemented by 李鲁鲁"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LFZTFVhzXgfv"
      },
      "source": [
        "检查下列变量是否被正确复制\n",
        "\n",
        "- [ ] openai的key（不要把自己的key 上传到github！）\n",
        "- [ ] 指定下载小说 wget后面的网址\n",
        "- [ ] 指定保存位置 save_folder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "krApyCoH3oJq"
      },
      "outputs": [],
      "source": [
        "\n",
        "import os\n",
        "\n",
        "key = 'sk-VvF4'\n",
        "key_bytes = key.encode()\n",
        "os.environ[\"OPENAI_API_KEY\"] = key_bytes.decode('utf-8')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kmb1oQFs3xsp",
        "outputId": "6d1f9ac0-3e75-44bc-c7d7-7d8741365f1e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2023-09-01 01:48:21--  https://raw.githubusercontent.com/LC1332/Prophet-Andrew-Ng/main/langchain/%E5%B0%84%E9%9B%95%E8%8B%B1%E9%9B%84%E4%BC%A0.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.110.133, 185.199.111.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2795446 (2.7M) [text/plain]\n",
            "Saving to: ‘射雕英雄传.txt’\n",
            "\n",
            "射雕英雄传.txt      100%[===================>]   2.67M  --.-KB/s    in 0.05s   \n",
            "\n",
            "2023-09-01 01:48:22 (56.2 MB/s) - ‘射雕英雄传.txt’ saved [2795446/2795446]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "!wget https://raw.githubusercontent.com/LC1332/Prophet-Andrew-Ng/main/langchain/%E5%B0%84%E9%9B%95%E8%8B%B1%E9%9B%84%E4%BC%A0.txt\n",
        "\n",
        "input_name = '/content/射雕英雄传.txt'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zlrE4zpzXmgd",
        "outputId": "8f5f0dcb-140f-4c51-eccd-2046eeaccc71"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "# mount google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "import os\n",
        "\n",
        "save_folder = \"/content/drive/MyDrive/StoryGPT/shediaoyingxiongzhuan_extract\"\n",
        "\n",
        "if not os.path.exists(save_folder):\n",
        "    os.makedirs(save_folder)\n",
        "else:\n",
        "    print('注意，文件夹',save_folder,'已经存在')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58Arx3x788Ko"
      },
      "source": [
        "# 下载和切片天龙八部"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "18SGBCpkmJcN"
      },
      "source": [
        "这里你要给小说切一下章节（最好）\n",
        "\n",
        "如果你非常非常懒\n",
        "\n",
        "可以使用\n",
        "\n",
        "```python\n",
        "chapters = [raw_text]\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oOIeYDaT4s_4"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "\n",
        "raw_text = open(input_name, encoding='utf-8').read()\n",
        "\n",
        "# chapter_pattern = r'第\\d+章[^\\n]+\\n'\n",
        "\n",
        "# chapters = raw_text.split(chapter_pattern)\n",
        "\n",
        "# print(len(chapters))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAgEOI23mTgv"
      },
      "source": [
        "我觉得这段代码现在比正则表达式好用\n",
        "\n",
        "当然这个代码是针对 第xx章 xxxx 这种格式的\n",
        "\n",
        "如果有的时候你是 ---第xx章--\n",
        "\n",
        "你要改一下那个startswith那个 或者把前面的特殊符号给strip掉"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nggnYZtpaizy",
        "outputId": "6b8f11f0-c4bb-4e8d-90a1-d275123cac44"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "40\n"
          ]
        }
      ],
      "source": [
        "chapters = []\n",
        "chapter_contents = []\n",
        "\n",
        "for line in raw_text.split('\\n'):\n",
        "    Flag = False\n",
        "    if line.strip().startswith('第'):\n",
        "        # 遇到章节标题,将之前章节内容添加到结果列表\n",
        "\n",
        "        head = line.strip()\n",
        "        # print(head)\n",
        "        head = head[:min(10,len(head))]\n",
        "        if head.find('章',1)>0:\n",
        "            # print(head)\n",
        "            Flag = True\n",
        "\n",
        "    if Flag:\n",
        "        if chapter_contents:\n",
        "            chapters.append('\\n'.join(chapter_contents))\n",
        "            chapter_contents = []\n",
        "        # 记录当前章节标题\n",
        "        # chapters.append(line)\n",
        "    else:\n",
        "        # 累积章节内容\n",
        "        chapter_contents.append(line)\n",
        "\n",
        "# 添加最后一个章节内容\n",
        "if chapter_contents:\n",
        "    chapters.append('\\n'.join(chapter_contents))\n",
        "\n",
        "print(len(chapters))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_wBzIPjhmhAc"
      },
      "source": [
        "看一看提取出来的章节对不对"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C2--yWclatnp",
        "outputId": "1d366f4e-178e-402a-942a-0daf946b6e44"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "0 \n",
            "钱塘江浩浩江水，日\n",
            "1 \n",
            "颜烈跨出房门，只见\n",
            "2 \n",
            "寺里僧众见焦木圆寂\n",
            "3 \n",
            "完颜洪熙笑道：“好\n",
            "4 \n",
            "一行人下得山来，走\n",
            "5 \n",
            "帐中六怪低声计议。\n",
            "6 \n",
            "江南六怪与郭靖晓行\n",
            "7 \n",
            "王处一脚步好快，不\n",
            "8 \n",
            "郭黄二人来到赵王府\n",
            "9 \n",
            "完颜康斗然见到杨铁\n",
            "10 \n",
            "沙通天见师弟危殆，\n",
            "11 \n",
            "黄蓉正要将鸡撕开，\n",
            "12 \n",
            "黄蓉回到客店安睡，\n",
            "13 \n",
            "只见五男一女，走进\n",
            "14 \n",
            "陆冠英扶起完颜康，\n",
            "15 \n",
            "郭、黄二人自程府出\n",
            "16 \n",
            "周伯通道：“你道是\n",
            "17 \n",
            "郭靖循着蛇声走去，\n",
            "18 \n",
            "洪七公万万想不到这\n",
            "19 \n",
            "洪七公与郭靖见欧阳\n",
            "20 \n",
            "欧阳锋只感身上炙热\n",
            "21 \n",
            "黄蓉见欧阳锋拖泥带\n",
            "22 \n",
            "洪七公、周伯通、郭\n",
            "23 \n",
            "黄蓉向外走了两步，\n",
            "24 \n",
            "黄药师仰天一笑，说\n",
            "25 \n",
            "黄药师心想不明不白\n",
            "26 \n",
            "两人正闹间，楼梯声\n",
            "27 \n",
            "此时鲁有脚已经醒转\n",
            "28 \n",
            "郭靖在雕背连声呼叫\n",
            "29 \n",
            "两人顺着山路向前走\n",
            "30 \n",
            "一灯大师低低叹了口\n",
            "31 \n",
            "穆念慈右手让黄蓉握\n",
            "32 \n",
            "郭靖与黄蓉此刻心意\n",
            "33 \n",
            "黄蓉飞舞正急，忽然\n",
            "34 \n",
            "船靠岸边，走上二三\n",
            "35 \n",
            "黄蓉幽幽的道：“欧\n",
            "36 \n",
            "这一日郭靖驻军那密\n",
            "37 \n",
            "郭靖解下长衣，执住\n",
            "38 \n",
            "郭靖纵马急驰数日，\n",
            "39 \n",
            "欧阳锋冷冷的道：“\n"
          ]
        }
      ],
      "source": [
        "for i,ch in enumerate(chapters):\n",
        "    print(i,ch[:10])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7lYbFWh2h44T",
        "outputId": "951eb7cf-d75c-4a6d-ae03-74f5adcf6a8f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "40\n",
            "\n",
            "钱塘江浩浩江水，日日夜夜无穷无休的从临安牛家村边绕过，东流入海。\n",
            "\n",
            "    江畔一排数十株乌柏树，叶子似火烧般红，正是八月天时。村前村后的野草刚起始变黄，一抹斜阳映照之下，更增了几分萧索。两株大松树\n"
          ]
        }
      ],
      "source": [
        "print(len(chapters))\n",
        "print(chapters[0][:100])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2GeitXNE69hg"
      },
      "outputs": [],
      "source": [
        "import tiktoken\n",
        "\n",
        "enc = tiktoken.get_encoding(\"cl100k_base\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LiG5bOg76-_-"
      },
      "outputs": [],
      "source": [
        "#@title  定义divide函数，用来切分超长文本\n",
        "def divide_str(s, sep=['\\n', '.', '。']):\n",
        "    mid_len = len(s) // 2  # 中心点位置\n",
        "    best_sep_pos = len(s) + 1  # 最接近中心点的分隔符位置\n",
        "    best_sep = None  # 最接近中心点的分隔符\n",
        "    for curr_sep in sep:\n",
        "        sep_pos = s.rfind(curr_sep, 0, mid_len)  # 从中心点往左找分隔符\n",
        "        if sep_pos > 0 and abs(sep_pos - mid_len) < abs(best_sep_pos -\n",
        "                                                        mid_len):\n",
        "            best_sep_pos = sep_pos\n",
        "            best_sep = curr_sep\n",
        "    if not best_sep:  # 没有找到分隔符\n",
        "        return s, ''\n",
        "    return s[:best_sep_pos + 1], s[best_sep_pos + 1:]\n",
        "\n",
        "\n",
        "def strong_divide(s):\n",
        "    left, right = divide_str(s)\n",
        "\n",
        "    if right != '':\n",
        "        return left, right\n",
        "\n",
        "    whole_sep = ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                 '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                 '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                 '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "               '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                 '·', '…']\n",
        "    left, right = divide_str(s, sep=whole_sep)\n",
        "\n",
        "    if right != '':\n",
        "        return left, right\n",
        "\n",
        "    mid_len = len(s) // 2\n",
        "    return s[:mid_len], s[mid_len:]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-cmOV5a8mlNW"
      },
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xbvld6mG7A-J",
        "outputId": "8bb96998-17b4-4afb-9973-0adec7be05d0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "929\n"
          ]
        }
      ],
      "source": [
        "#@title 以1500 token为限，切分chunk，输出总chunk数量\n",
        "\n",
        "max_token_len = 1500\n",
        "chunk_text = []\n",
        "\n",
        "\n",
        "for chapter in chapters:\n",
        "\n",
        "    split_text = chapter.split('\\n')\n",
        "\n",
        "    curr_len = 0\n",
        "    curr_chunk = ''\n",
        "\n",
        "    tmp = []\n",
        "\n",
        "    for line in split_text:\n",
        "        line_len = len(enc.encode( line ))\n",
        "\n",
        "        if line_len <= max_token_len - 5:\n",
        "            tmp.append(line)\n",
        "        else:\n",
        "            # print('divide line with length = ', line_len)\n",
        "            path = [line]\n",
        "            tmp_res = []\n",
        "\n",
        "            while path:\n",
        "                my_str = path.pop()\n",
        "                left, right = strong_divide(my_str)\n",
        "\n",
        "                len_left = len(enc.encode( left ))\n",
        "                len_right = len(enc.encode( right ))\n",
        "\n",
        "                if len_left > max_token_len - 15:\n",
        "                    path.append(left)\n",
        "                else:\n",
        "                    tmp_res.append(left)\n",
        "\n",
        "                if len_right > max_token_len - 15:\n",
        "                    path.append(right)\n",
        "                else:\n",
        "                    tmp_res.append(right)\n",
        "\n",
        "            for line in tmp_res:\n",
        "                tmp.append(line)\n",
        "\n",
        "    split_text = tmp\n",
        "\n",
        "    for line in split_text:\n",
        "        line_len = len(enc.encode( line ))\n",
        "\n",
        "        if line_len > max_token_len:\n",
        "            print('warning line_len = ', line_len)\n",
        "\n",
        "        if curr_len + line_len <= max_token_len:\n",
        "            curr_chunk += line\n",
        "            curr_chunk += '\\n'\n",
        "            curr_len += line_len\n",
        "            curr_len += 1\n",
        "        else:\n",
        "            chunk_text.append(curr_chunk)\n",
        "            curr_chunk = line\n",
        "            curr_len = line_len\n",
        "\n",
        "    if curr_chunk:\n",
        "        chunk_text.append(curr_chunk)\n",
        "\n",
        "    # break\n",
        "\n",
        "print(len(chunk_text))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XzB5ZszL7rcr"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OCKTCEB29DlH"
      },
      "source": [
        "# 编辑新的Kor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sNXooBMC9Hdz"
      },
      "source": [
        "新的Kor要简单很多，我们准备不再提取，动作。只保留对话\n",
        "\n",
        "这里的Kor建议你调一调 当然如果很懒，也可以不调了"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0S6ZPyOK9F_N"
      },
      "outputs": [],
      "source": [
        "from kor.extraction import create_extraction_chain\n",
        "from kor.nodes import Object, Text, Number\n",
        "from langchain.chat_models import ChatOpenAI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CGKDYtmE-wiI"
      },
      "outputs": [],
      "source": [
        "from langchain.llms import OpenAI\n",
        "llm = ChatOpenAI(\n",
        "    model_name=\"gpt-3.5-turbo\",\n",
        "    temperature=0\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KYmO_rs2FMW0",
        "outputId": "3685de6f-7631-405f-92e8-779214778b82"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "    张翠山沉吟道：“这梅花小镖用左手发射？少林派门下怎地出现了道人，莫非也是乔装的？”那少女微笑道：“道士扮和尚须剃光头，和尚扮道士却容易得多，戴顶道冠便成。”张翠山点了点头。那少女道：“我心知此事不妙，但那瘦子我尚自抵敌不过，那道人似乎更厉害得多，何况他们共有六人？这可没了计较。”张翠山张口欲言，但终于忍住了。那少女道：“我猜你是想问：‘干么不上武当山来跟我们说明？’是不是？我可不能上武当山啊，倘若我自己能出面，又何必委托都大锦走这趟镖呢？我徬徨无计，在道上闷走，恰好撞到你跟都大锦他们说话。后来见你去找寻俞三侠，我想武当七侠正主儿已接上了手，不用我再凑热闹，凭我这点微末本领，也帮不了甚么忙。那时我急于解毒，便即东还，不知俞三侠后来怎样了？”张翠山当下说了俞岱岩受人毒害的情状。那少女长叹一声，睫毛微微颤动，说道：“但愿俞三侠吉人天相，终能治愈，否则……否则……”张翠山听她语气诚恳，心下感激，说道：“多谢姑娘好心。”说着眼眶微湿。那少女摇了摇头，说道：“我回到江南，叫人一看这梅花镖，有人识得是少林派的独门暗器，说道除非是发暗器之人的本门解药，否则毒性难除。临安府除了龙门镖局，还有谁是少林派？于是我夜入镖局，要逼他们给解药，岂知他们不但不给，还埋伏下了人马，我一进门便对我猛下毒手。”张翠山“嗯”了一声，沉吟道：“你说故意安排，教他们认作是我？”那少女脸有腼腆之色，低下了头，轻轻的道：“我见你到衣铺去买了这套衣巾，觉得穿戴起来很是……很是好看，于是我跟着也头了一套。”张翠山道：“这便是了。只是你一出手便连杀数十人，未免过于狠辣，镖局中的人跟你又没怨仇。”那少女沉下脸来，冷笑道：“你要教训我么？我活了一十九岁，倒还没听人教训过呢。张五侠大仁大义，这就请罢。我这般心狠手辣之辈，原没盼望跟你结交。”\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(chunk_text[100])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "17_FRAsc-6Mc"
      },
      "outputs": [],
      "source": [
        "#@title 仅仅抽取dialogue的部分\n",
        "\n",
        "schema = Object(\n",
        "    id=\"script\",\n",
        "    description=\"Extract Dialogue in order From Novel, ignore the non-dialogue parts\",\n",
        "    attributes=[\n",
        "        Text(\n",
        "            id=\"role\",\n",
        "            description=\"The character who is speaking, use context to predict the name of the role.\",\n",
        "        ),\n",
        "        Text(\n",
        "            id=\"dialogue\",\n",
        "            description=\"The dialogue spoken by the characters in the sentence\",\n",
        "        ),\n",
        "    ],\n",
        "    examples=[\n",
        "        (\n",
        "            '''``村民中走出一个二十来岁的人汉，说道：“张先生，你可是从北方来吗”\n",
        "\n",
        "    张十五见他身材魁梧，浓眉大眼，便道：“正是。”那大汉道：“小弟作东，请先生去饮上三杯如何”张十五大喜，说道：“素不相识，怎敢叨扰”\n",
        "\n",
        "    那大汉笑道：“喝上三怀，那便相识了。我姓郭，名叫郭啸天。”``''',\n",
        "            [\n",
        "                {\"role\": \"郭啸天\", \"dialogue\": \"张先生，你可是从北方来吗\"},\n",
        "                {\"role\": \"张十五\",\"dialogue\": \"正是。\"},\n",
        "                {\"role\": \"郭啸天\", \"dialogue\": \"小弟作东，请先生去饮上三杯如何\"},\n",
        "                {\"role\": \"张十五\",\"dialogue\": \"素不相识，怎敢叨扰\"},\n",
        "                {\"role\": \"郭啸天\", \"dialogue\": \"喝上三怀，那便相识了。我姓郭，名叫郭啸天。\"}\n",
        "            ],\n",
        "        )\n",
        "    ],\n",
        "    many=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-srQ6KcPBuVz",
        "outputId": "97a8a482-ced6-4bd5-b05f-a73a10516d15"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Your goal is to extract structured information from the user's input that matches the form described below. When extracting information please make sure it matches the type information exactly. Do not add any attributes that do not appear in the schema shown below.\n",
            "\n",
            "```TypeScript\n",
            "\n",
            "script: Array<{ // Extract Dialogue in order From Novel, ignore the non-dialogue parts\n",
            " role: string // The character who is speaking, use context to predict the name of the role.\n",
            " dialogue: string // The dialogue spoken by the characters in the sentence\n",
            "}>\n",
            "```\n",
            "\n",
            "\n",
            "Please output the extracted information in CSV format in Excel dialect. Please use a | as the delimiter. \n",
            " Do NOT add any clarifying information. Output MUST follow the schema above. Do NOT add any additional columns that do not appear in the schema.\n",
            "\n",
            "\n",
            "\n",
            "Input: ``村民中走出一个二十来岁的人汉，说道：“张先生，你可是从北方来吗”\n",
            "\n",
            "    张十五见他身材魁梧，浓眉大眼，便道：“正是。”那大汉道：“小弟作东，请先生去饮上三杯如何”张十五大喜，说道：“素不相识，怎敢叨扰”\n",
            "\n",
            "    那大汉笑道：“喝上三怀，那便相识了。我姓郭，名叫郭啸天。”``\n",
            "Output: role|dialogue\n",
            "郭啸天|张先生，你可是从北方来吗\n",
            "张十五|正是。\n",
            "郭啸天|小弟作东，请先生去饮上三杯如何\n",
            "张十五|素不相识，怎敢叨扰\n",
            "郭啸天|喝上三怀，那便相识了。我姓郭，名叫郭啸天。\n",
            "\n",
            "Input: [user input]\n",
            "Output:\n"
          ]
        }
      ],
      "source": [
        "chain = create_extraction_chain(llm, schema)\n",
        "print(chain.prompt.format_prompt(text=\"[user input]\").to_string())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dGCRURq0CN5T"
      },
      "outputs": [],
      "source": [
        "i = 250\n",
        "response = chain.run( f\"``{chunk_text[i]}``\" )[\"data\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "v333VST6imbH",
        "outputId": "039cb05e-0d1e-45bd-ad43-b62d1595a6ee"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'script': [{'role': '王处一', 'dialogue': '你武功可比你爹爹强得多呀，那是怎么回事'}, {'role': '穆念慈', 'dialogue': '晚辈十三岁那年，曾遇到一位异人。他指点了我三天武功，可惜我生性愚鲁，没能学到甚么。'}, {'role': '王处一', 'dialogue': '他只教你三天，你就能胜过你爹爹。这位高人是谁'}, {'role': '穆念慈', 'dialogue': '不是晚辈胆敢隐瞒道长，实是我曾立过誓，不能说他的名号。'}, {'role': '王处一', 'dialogue': '丘师哥，你教完颜康教了有八九年吧'}, {'role': '丘处机', 'dialogue': '整整九年零六个月，唉，想不到这小子如此混蛋。'}, {'role': '王处一', 'dialogue': '这倒奇了'}, {'role': '丘处机', 'dialogue': '怎么'}, {'role': '柯镇恶', 'dialogue': '丘道长，你怎么找到杨大哥的后裔'}, {'role': '丘处机', 'dialogue': '说来也真凑巧。自从贫道和各位订了约会之后，到处探访郭杨两家的消息，数年之中，音讯全无，但总不死心，这年又到临安府牛家村去查访，恰好见到有几名公差到杨大哥的旧居来搬东西。贫道跟在他们背后，偷听他们说话，这几个人来头不小，竟是大金国赵王府的亲兵，奉命专程来取杨家旧居中一切家私物品，说是破凳烂椅，铁枪犁头，一件不许缺少。'}, {'role': '丘处机', 'dialogue': '贫道晚上夜探王府，要瞧瞧赵王万里迢迢的搬运这些破烂物事，到底是何用意。一探之后，不禁又是气愤，又是难受，原来杨兄弟的妻子包氏已贵为王妃。贫道大怒之下，本待将她一剑杀却，却见她居于砖房小屋之中，抚摸杨兄弟铁枪，终夜哀哭；心想她倒也不忘故夫，并非全无情义，这才饶了她性命。后来查知那小王子原来是杨兄弟的骨血，隔了数年，待他年纪稍长，贫道就起始传他武艺。'}, {'role': '柯镇恶', 'dialogue': '那小子是一直不知自己的身世的了'}, {'role': '丘处机', 'dialogue': '贫道也曾试过他几次口风，见他贪恋富贵，不是性情中人，是以始终不曾点被。几次教诲他为人立身之道，这小子只是油腔滑调的对我敷衍。若不是和七位有约，贫道哪有这耐心跟他穷耗本待让他与郭家小世兄较艺之后，不论谁胜谁败，咱们双方和好，然后对那小子说明他的身世，接他母亲出来，择地隐居。岂料杨兄弟尚在人世，而贫道和马师哥两人又着了好人暗算，终究救不得杨兄弟夫妇的性命，唉'}]}\n"
          ]
        }
      ],
      "source": [
        "print(response)\n",
        "# print(chunk_text[i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PWFoffIQDTi2"
      },
      "source": [
        "理论上这个抽取出来是这样的\n",
        "\n",
        "```\n",
        "role|dialogue\n",
        "干光豪|自今而后,咱二人再也不分什么东宗西宗啦。我俩东宗西宗联姻,合为一体......\n",
        "葛师妹|别......别这样。\n",
        "干光豪|你依了我,若是我日后负心,就掉在这水里,变个大忘八。\n",
        "葛师妹|你做忘八,可不是骂我不规矩吗?\n",
        "干光豪|什么人?\n",
        "干光豪|葛师妹,你拦住了那边山口!\n",
        "干光豪|前面是本派禁地,任何外人不得擅入。你再向前数丈,干犯禁忌,可叫你死葬身之地。\n",
        "段誉|我要性命,这才逃走......\n",
        "段誉|啊哟!\n",
        "```\n",
        "\n",
        "但是我遇到了latin error 所以下次在继续调试这部分\n",
        "\n",
        "让我们来看看摘要的部分"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F0EawLOADbnj"
      },
      "source": [
        "# Bullet形式摘要抽取"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ANBKqeO6CXFP"
      },
      "outputs": [],
      "source": [
        "system_prompt = \"\"\"\n",
        "Summarize the key points of the following text in a concise way, using bullet points.\n",
        "\"\"\"\n",
        "\n",
        "q_example = \"\"\"###\n",
        "Text:\n",
        "洪七公、周伯通、郭靖、黄蓉四人乘了小船，向西驶往陆地。郭靖坐在船尾扳桨，黄蓉不住向周伯通详问骑鲨游海之事，周伯通兴起，当场就要设法捕捉鲨鱼，与黄蓉大玩一场。\n",
        "郭靖见师父脸色不对，问道：“你老人家觉得怎样”洪七公不答，气喘连连，声息粗重。他被欧阳锋以“透骨打穴法”点中之后，穴道虽已解开，内伤却又加深了一层。黄蓉喂他服了几颗九花玉露丸，痛楚稍减，气喘仍是甚急。\n",
        "老顽童不顾别人死活，仍是嚷着要下海捉鱼，黄蓉却已知不妥，向他连使眼色，要他安安静静的，别吵得洪七公心烦。周伯通并不理会，只闹个不休。黄蓉皱眉道：“你要捉鲨鱼，又没饵引得鱼来，吵些甚么”\n",
        "\n",
        "Summarize in BULLET POINTS form:\n",
        "\"\"\"\n",
        "\n",
        "a_example = \"\"\"\n",
        "- 洪七公等四人乘船西行,洪七公因受内伤加重而气喘不止\n",
        "- 周伯通要捉鲨鱼玩,被黄蓉阻止以免掀翻小船\n",
        "\"\"\"\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x9oE6AYhGTCu"
      },
      "outputs": [],
      "source": [
        "from langchain.chat_models import ChatOpenAI\n",
        "from langchain.prompts.chat import (\n",
        "    ChatPromptTemplate,\n",
        "    SystemMessagePromptTemplate,\n",
        "    AIMessagePromptTemplate,\n",
        "    HumanMessagePromptTemplate,\n",
        ")\n",
        "from langchain.schema import (\n",
        "    AIMessage,\n",
        "    HumanMessage,\n",
        "    SystemMessage\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VS-1K7mRGTbR"
      },
      "outputs": [],
      "source": [
        "llm = ChatOpenAI(\n",
        "    model_name=\"gpt-3.5-turbo\",\n",
        "    temperature=0\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ar_nQWRoGXJL",
        "outputId": "a155c50f-6196-4bb7-e290-6c220de93080"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "- 完颜康和陆冠英进行比武较量\n",
            "- 陆冠英使用九阴白骨爪和全真派拳法对抗完颜康的手爪功夫\n",
            "- 陆冠英展示了他的外家功夫，特别是腿上的技巧\n",
            "- 完颜康在战斗中受伤，但最终使用绝招“怀心腿”将陆冠英击败\n",
            "- 陆冠英被完颜康推倒，撞向陆庄主，引发了一场争斗\n"
          ]
        }
      ],
      "source": [
        "messages = [SystemMessage( content = system_prompt),\n",
        "            HumanMessage( content = q_example),\n",
        "            AIMessage( content = a_example)]\n",
        "\n",
        "sel_id = 300\n",
        "\n",
        "new_input = f\"\"\"###\n",
        "Text:\n",
        "{chunk_text[sel_id]}\n",
        "\n",
        "Summarize in BULLET POINTS form:\"\"\"\n",
        "\n",
        "messages.append( HumanMessage(content = new_input) )\n",
        "\n",
        "response = llm( messages )\n",
        "\n",
        "print(response.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TfJ0Fnl_Qyhr"
      },
      "source": [
        "理论上输出是这样的\n",
        "\n",
        "\n",
        "- 段誉和朱丹臣讨论诗词,表达友谊\n",
        "- 木婉清不懂诗词,觉得两人说书呆子话\n",
        "- 朱丹臣礼貌对木婉清,想请两人回段家\n",
        "- 段誉担心鳄神,同意回家,让朱丹臣护送\n",
        "- 木婉清说朱丹臣话太多,自己是山野女子不懂\n",
        "- 朱丹臣表示会简化自己的文词,不让木婉清觉得难懂\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WfA6NoYNQxXZ",
        "outputId": "cb63a0d5-da20-45d1-e0fc-4c562acafa23"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/content/drive/MyDrive/StoryGPT/shediaoyingxiongzhuan_extract\n"
          ]
        }
      ],
      "source": [
        "print(save_folder)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SM63hA11HQTa",
        "outputId": "71fe217a-6ff1-43a4-fc60-6536c75839e5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1NLzjtsZOHiO"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hCW42kX9O5pO",
        "outputId": "0f6064a7-ddba-450f-b904-1596403764a6"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "  6%|▌         | 57/929 [24:58<7:11:13, 29.67s/it]"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import json\n",
        "from tqdm import tqdm\n",
        "\n",
        "# save_folder = \"/content/drive/MyDrive/GPTData/weixiaobao_extract\"\n",
        "\n",
        "for i in tqdm(range(len(chunk_text))):\n",
        "    save_name = os.path.join(save_folder, f\"{i}_dialogue.txt\")\n",
        "\n",
        "    if not os.path.exists(save_name) or os.path.getsize(save_name) < 5:\n",
        "        if os.path.exists(save_name):\n",
        "            print('re-generate dialogue id = ', i)\n",
        "        query_text = f\"``{chunk_text[i]}``\"\n",
        "        dialogue_response = chain.run( query_text )[\"data\"]\n",
        "\n",
        "        with open(save_name, 'w', encoding='utf-8') as f:\n",
        "            if 'script' not in dialogue_response:\n",
        "                print('Error: response does not contain key \"script\"')\n",
        "            else:\n",
        "                for chat in dialogue_response['script']:\n",
        "                    json_str = json.dumps(chat, ensure_ascii=False)\n",
        "                    f.write(json_str+\"\\n\")\n",
        "\n",
        "    save_name_sum = os.path.join(save_folder, f\"{i}_sum.txt\")\n",
        "\n",
        "    if not os.path.exists(save_name_sum) or os.path.getsize(save_name_sum) < 5:\n",
        "        if os.path.exists(save_name_sum):\n",
        "            print('re-summarize id = ',i )\n",
        "        #dealing with summarize\n",
        "        messages = [SystemMessage( content = system_prompt),\n",
        "                HumanMessage( content = q_example),\n",
        "                AIMessage( content = a_example)]\n",
        "\n",
        "        new_input = f\"###\\nText:\\n{chunk_text[ i ]}\\nSummarize in BULLET POINTS form:\"\n",
        "\n",
        "        messages.append( HumanMessage(content = new_input) )\n",
        "\n",
        "        summarize_response = llm( messages ).content\n",
        "\n",
        "        with open(save_name_sum, 'w', encoding='utf-8') as f:\n",
        "            f.write( summarize_response )\n",
        "\n",
        "    raw_text_save_name = os.path.join(save_folder, f\"{i}_raw.txt\")\n",
        "\n",
        "    if not os.path.exists(raw_text_save_name) or os.path.getsize(raw_text_save_name) < 5:\n",
        "        with open(raw_text_save_name, 'w', encoding='utf-8') as f:\n",
        "            f.write( chunk_text[i] )\n",
        "\n",
        "    # if i >5:\n",
        "    #     break\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gbSAvk4mivfw",
        "outputId": "1b4ff050-58f7-4bf1-9697-f0cca809cb52"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "    游坦之乘着众人扰攘之际，和全冠清低声商议了一阵，又朗声道：“我大宋国步艰危，江湖同道却又不能齐心合力，以至时受番帮欺压。因此丐帮主张立一位武林盟主，大伙儿听奉号令，有什么大事发生，便不致乱成一团了。玄慈方丈，你赞不赞成？”    玄慈缓缓地道：“庄帮主的话，倒也言之成理。但老衲有一事不解，却要请教。”游坦之道：“什么事？”玄慈道：“庄帮主已拜丁先生为师，算是星宿派门人了，是也不是？”游坦之道：“这个……这是我自己的事，与你无关。”玄慈道：“星宿派乃西域门派，非我大宋武林同道。我大宋立不立武林盟主，可与星宿派无涉。就算中原武林同道要推举一位盟主，以便统筹事功，阁下是星宿派门人，却也不便参与了。”\n",
            "    众英雄纷纷说道：“不错！”“少林方丈之言甚是。”“你是番邦门派的走狗奴才，怎可妄想做我中原武林的盟主？。”\n",
            "    游坦之无言可答，向丁春秋望望，又向全冠清瞧瞧，盼望他们出言解围。\n",
            "    丁春秋咳嗽一声，说道：“少林方丈言之差矣！老夫乃山东曲阜人氏，生于圣人之邦，星宿派乃老夫一手创建，怎能说是西域番邦的门派？星宿派虽居处西域，那只不过是夫夫暂时隐居之地。你说星宿派是番邦门派，那么孔夫子也是番邦人氏了，可笑啊可笑！说到西域番邦，少林武功源于天竺达摩祖师，连佛教也是西域番邦之物，我看少林派才是西域的门派呢！”此言一出，玄慈和群雄都感不易抗辩。\n",
            "    全冠清朗声道：“天下武功，源流难考。西域武功传于中土者有之，中土武功传于西域者亦有之。我帮庄帮主乃中土人氏，丐帮素为中原门派，他自然是中原武林的领袖人物。玄慈方丈，今日之事，当以武功强弱定胜负，不以言辞舌辩定输赢。丐帮与少林派到底谁强谁弱，只须你们两位首领出手较量，高下立判，否则便是说上半天，又有何益？倘若你有自知之明，不是敝帮庄帮主的敌手，只须甘拜下风，推戴我庄帮主为武林盟主，倒也不是非出手不可的。”这几句话，显然认定玄慈是明知不敌，胆怯推诿。\n",
            "    玄慈向前走了几步，说道：“庄帮主，你既非要老衲出手不可，老衲若再顾念贵帮和敝帮数百年的交情，坚不肯允，倒是对贵帮不敬了。”眼光向群雄缓缓掠过，朗声道：“天下英雄，今日人人亲眼目睹，我少林派决无与丐帮争雄斗胜之意，实是丐帮帮主步步见逼，老衲退无可退，避无可避。”\n",
            "    群雄纷纷说道：“不错，咱们都是见证，少林派并无丝毫理亏之处。”\n",
            "    游坦之只是挂念着阿紫的安危，一心要尽快杀了玄慈，好向丁春秋交差，大声说道：“比武较量，强存弱亡，说不上谁理亏不理亏，快快上来动手吧！”\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(chunk_text[984])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XmB_H2Z4CGLr"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyPNrUOAmqr1/Zv6rqc/8ly5",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
